import mindspore.nn as nn
from mindspore.ops import operations as P
import mindspore as ms



class ClassificationHead(nn.Cell):
    def __init__(self, network, feature_dim, num_classes):
        super().__init__()
        self.network = network
        self.num_classes = num_classes
        self.head = nn.Dense(feature_dim, num_classes)
        self.loss = nn.CrossEntropyLoss()
    
    def construct(self, image, label=None):
        x = self.head(self.network(image))
        if self._phase == 'train':
            label = label.astype(ms.int32)
            l = self.loss(x, label)
            return l
        else:
            return x


